using System;
using System.Collections.Generic;
using System.Linq;

using Atomic.Libraries.Mathematics;
using Atomic.Structures;
using Atomic.Thermodynamics;
using Atomic.Thermodynamics.PartitionFunctions;

namespace Atomic.Thermodynamics.PartitionFunctions
{
	[Serializable]
	public class PartitionFunctionOLD
	{
		private int n;
		private Atom atom;
		private double[] counts; // use absolute per-unit cell particle numbers instead! (as energies are absolute per-unit cell values)
		private IHelmholtzPotential[] potentials;
		private int[] degeneracies;
		private double kb, t, beta;

		public PartitionFunctionOLD(IEnumerable<IPartitionFunctionEntry> entries)
		{
			if (entries == null)
			{
				throw new ArgumentNullException();
			}

			n = entries.Count();

			if (n == 0)
			{
				throw new ArgumentException();
			}

			Atom[] atoms = entries.SelectMany(e => e.Composition.Atoms).Distinct().OrderBy(a => a == Atom.Vacancy).ToArray();
			potentials = entries.Select(e => e.Potential).ToArray();
			degeneracies = entries.Select(e => e.Degeneracy).ToArray();

			if (atoms.Length != 2 || atoms[1] != Atom.Vacancy)
			{
				// Ordered with the vacancy as the last atom above (false is ordered before true).
				throw new ArgumentException("Currently only binary systems are supported. One of the atoms need to by a vacancy.");
			}

			// The non-vacancy atom and corresponding per-unit cell counts.
			atom = atoms[0];
			counts = entries.Select(p => p.Composition.Count(atom)).ToArray();

			kb = 8.617332478e-5; // eV/K
			t = 723.0; // K
			beta = 1.0 / (kb * t);
		}

		private void LogPartitionFunction(double v, double mu, out double lz, out double lz_v, out double lz_mu, out double[] f, out double[] f_v, out double[] f_mu)
		{
			double[] e = new double[n];
			double[] e_v = new double[n];
			double[] e_mu = new double[n];
			double e0 = double.NegativeInfinity;
			double e0_v = double.NaN;
			double e0_mu = double.NaN;
			for (int i = 0; i < n; i++)
			{
				e[i] = beta * mu * counts[i] - beta * potentials[i].FreeEnergy(t, v);
				e_v[i] = beta * potentials[i].Pressure(t, v);
				e_mu[i] = beta * counts[i];

				//e0 = Math.Max(e0, e[i]);
				if (e[i] > e0)
				{
					e0 = e[i];
					e0_v = e_v[i];
					e0_mu = e_mu[i];
				}
			}

			// Compute partition function with prefactor introduced for numerical stability.
			double z = 0.0;
			double z_v = 0.0;
			double z_mu = 0.0;
			f = new double[n];
			f_v = new double[n];
			f_mu = new double[n];
			for (int i = 0; i < n; i++)
			{
				f[i] = Math.Exp(e[i] - e0) * degeneracies[i];
				f_v[i] = f[i] * (e_v[i] - e0_v);
				f_mu[i] = f[i] * (e_mu[i] - e0_mu);
				z += f[i];
				z_v += f_v[i];
				z_mu += f_mu[i];
			}

			for (int i = 0; i < n; i++)
			{
				// Prefactor doesn't matter for the Boltzmann factor.
				f_v[i] = f_v[i] / z - f[i] / (z * z) * z_v;
				f_mu[i] = f_mu[i] / z - f[i] / (z * z) * z_mu;
				f[i] /= z;
			}

			// Divide out prefactor again after taking the logarithm.
			lz = Math.Log(z) + e0;
			lz_v = 1.0 / z * z_v + e0_v;
			lz_mu = 1.0 / z * z_mu + e0_mu;
		}
		
		public double GrandPotential(double v, double mu)
		{
			double phi, phi_v, phi_mu;
			GrandPotential(v, mu, out phi, out phi_v, out phi_mu);

			return phi;
		}

		public double GrandPotentialDerivative(double v, double mu)
		{
			double phi, phi_v, phi_mu;
			GrandPotential(v, mu, out phi, out phi_v, out phi_mu);

			return phi_v;
		}
			
		public double GrandPotentialDerivativeNumerical(double v, double mu)
		{
			double dv = 0.0001;
			double phi_v = (GrandPotential(v + dv, mu) - GrandPotential(v, mu)) / dv;

			return phi_v;
		}

		public void GrandPotential(double v, double mu, out double phi, out double phi_v, out double phi_mu)
		{
			double lz, lz_v, lz_mu;
			double[] f, f_v, f_mu;
			LogPartitionFunction(v, mu, out lz, out lz_v, out lz_mu, out f, out f_v, out f_mu);

			phi = -lz / beta;
			phi_v = -lz_v / beta;
			phi_mu = -lz_mu / beta;
		}

		public double ParticleNumber(double v, double mu)
		{
			double x, x_v, x_mu;
			ParticleNumber(v, mu, out x, out x_v, out x_mu);

			return x;
		}	

		public double ParticleNumberDerivative(double v, double mu)
		{
			double x, x_v, x_mu;
			ParticleNumber(v, mu, out x, out x_v, out x_mu);

			return x_v;
		}

		public double ParticleNumberDerivativeNumerical(double v, double mu)
		{
			double dv = 0.0001;
			double x_v = (ParticleNumber(v + dv, mu) - ParticleNumber(v, mu)) / dv;

			return x_v;
		}

		public double ParticleNumberChemicalPotentialDerivative(double v, double mu)
		{
			double dmu = 0.0001;
			double x_mu = (ParticleNumber(v, mu + dmu) - ParticleNumber(v, mu)) / dmu;

			return x_mu;
		}

		public void ParticleNumber(double v, double mu, out double x, out double x_v, out double x_mu)
		{
			double lz, lz_v, lz_mu;
			double[] f, f_v, f_mu;
			LogPartitionFunction(v, mu, out lz, out lz_v, out lz_mu, out f, out f_v, out f_mu);

			x = 0.0;
			x_v = 0.0;
			x_mu = 0.0;
			for (int i = 0; i < n; i++)
			{
				// Prefactor doesn't matter for the Boltzmann factor.
				x += f[i] * counts[i];
				x_v += f_v[i] * counts[i];
				x_mu += f_mu[i] * counts[i];
			}
		}

		public double ChemicalPotential(double v, double x)
		{
			double mu, mu_v;
			ChemicalPotential(v, x, out mu, out mu_v);

			return mu;
		}
		
		public double ChemicalPotentialDerivative(double v, double x)
		{
			double mu, mu_v;
			ChemicalPotential(v, x, out mu, out mu_v);

			return mu_v;
		}

		public double ChemicalPotentialDerivativeNumerical(double v, double mu)
		{
			double dv = 0.0001;
			double mu_v = (ChemicalPotential(v + dv, mu) - ChemicalPotential(v, mu)) / dv;

			return mu_v;
		}

		public void ChemicalPotential(double v, double x, out double mu, out double mu_v)
		{
			mu = RootFinder.Bisection(y => ParticleNumber(v, y) - x, -1.0e8, 1.0e8);

			// Ignore the updated particle number.
			double x_v, x_mu;
			ParticleNumber(v, mu, out x, out x_v, out x_mu);

			// Compute derivative implicitly from the inverse of the Jacobian of the change-of-variable mapping $(V,\mu)\mapsto (V,N(V,\mu)$.
			mu_v = -x_v / x_mu;
		}

		public double GibbsPotentialChemicalPotentialNumerical(double p, double x)
		{
			double h = 0.0006;
			double xh = x + h;

			double f = GibbsPotential(p, x);
			double fh = GibbsPotential(p, xh);

			double mu = (fh - f) / h;

			return mu;
		}

		public double HelmholtzPotential(double v, double x)
		{
			double a, a_v;
			HelmholtzPotential(v, x, out a, out a_v);

			return a;
		}

		public double HelmholtzPotentialDerivative(double v, double x)
		{
			double a, a_v;
			HelmholtzPotential(v, x, out a, out a_v);

			return a_v;
		}

		public double HelmholtzPotentialDerivativeNumerical(double v, double x)
		{
			double dv = 0.0001;
			double a_v = (HelmholtzPotential(v + dv, x) - HelmholtzPotential(v, x)) / dv;

			return a_v;
		}

		public void HelmholtzPotential(double v, double x, out double a, out double a_v)
		{
			// Legendre transform of the grand-potential.
			double mu, mu_v;
			ChemicalPotential(v, x, out mu, out mu_v);

			double phi, phi_v, phi_mu;
			GrandPotential(v, mu, out phi, out phi_v, out phi_mu);

			a = phi + mu * x;
			a_v = phi_v + phi_mu * mu_v + mu_v * x;
		}

		public double Pressure(double v, double x)
		{
			double a, a_v;
			HelmholtzPotential(v, x, out a, out a_v);

			// Pressure is given as this derivative.
			double p = -a_v;

			return p;
		}

		public double PressureNumerical(double v, double x)
		{
			// FIXME Find closed form solution!!
			double dv = 0.0001;
			double a_v = (HelmholtzPotential(v + dv, x) - HelmholtzPotential(v, x)) / dv;
			return -a_v;
		}

		public double EquilibriumVolume(double p, double x)
		{
			return RootFinder.Bisection(v => Pressure(v, x) - p, 70.0, 101.0);
		}

		public double GibbsPotential(double p, double x)
		{
			// Legendre transform here again!
			double v = EquilibriumVolume(p, x);
			return HelmholtzPotential(v, x) + p * v;
		}
	}
}
